#include "cpu_kernel.h"
#include <cstring>   // For memcpy if needed, though C.fill might suffice
#include <algorithm> // For std::min, std::max
#include <iostream>  // For std::cout, std::cerr
#include <stdexcept> // For std::invalid_argument
#include <cblas.h>   // 包含 OpenBLAS C 接口头文件

// Basic serial matrix multiplication implementation (Keep this)
void matrixMultiplySerial(const Matrix& A, const Matrix& B, Matrix& C) {
    int M = A.getRows();
    int K = A.getCols();
    int N = B.getCols();
    
    if (B.getRows() != K || C.getRows() != M || C.getCols() != N) {
        throw std::invalid_argument("matrixMultiplySerial: Matrix dimensions do not match.");
    }
    
    C.fill(0.0f);
    
    const float* a_data = A.getData();
    const float* b_data = B.getData();
    float* c_data = C.getData();

    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += a_data[i * K + k] * b_data[k * N + j];
            }
            c_data[i * N + j] = sum;
        }
    }
}
    
// Optimized CPU matrix multiplication using OpenBLAS
void matrixMultiplyOptimized(const Matrix& A, const Matrix& B, Matrix& C) {
    int M = A.getRows();
    int K = A.getCols();
    int N = B.getCols();
    
    if (B.getRows() != K || C.getRows() != M || C.getCols() != N) {
        throw std::invalid_argument("matrixMultiplyOptimized: Matrix dimensions do not match.");
    }
    
    // 注意: 不再需要手动实现小矩阵切换或分块，OpenBLAS 内部会处理这些优化
    std::cout << "Info: Using OpenBLAS cblas_sgemm for CPU matrix multiplication." << std::endl;

    const float* a_data = A.getData();
    const float* b_data = B.getData();
    float* c_data = C.getData();

    // OpenBLAS sgemm 调用: C = alpha * A * B + beta * C
    // 参数解释:
    // CblasRowMajor: 指定矩阵是行主序存储 (我们的 Matrix 类是这样)
    // CblasNoTrans: 指定 A 矩阵不进行转置
    // CblasNoTrans: 指定 B 矩阵不进行转置
    // M: A 和 C 的行数
    // N: B 和 C 的列数
    // K: A 的列数 / B 的行数
    // 1.0f: alpha 值 (A*B 的缩放因子)
    // a_data: 指向 A 矩阵数据的指针
    // K: A 矩阵的主维度 (leading dimension, 行主序下是列数)
    // b_data: 指向 B 矩阵数据的指针
    // N: B 矩阵的主维度 (leading dimension, 行主序下是列数)
    // 0.0f: beta 值 (C 矩阵初始值的缩放因子, 0.0f 表示不保留 C 的初始值)
    // c_data: 指向 C 矩阵数据的指针
    // N: C 矩阵的主维度 (leading dimension, 行主序下是列数)
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, 
                M, N, K, 
                1.0f, a_data, K, 
                      b_data, N, 
                0.0f, c_data, N);

    // 注意: OpenBLAS 通常会根据环境变量 (如 OMP_NUM_THREADS) 或内部逻辑自动使用多线程。
    // 原来的 #pragma omp parallel for 块被替换为上面的单行调用。
}
